[PyTorch] Advance FusedAdam step counter for empty param groups - #3318
[PyTorch] Advance FusedAdam step counter for empty param groups#3318adityasingh2400 wants to merge 1 commit into
Conversation
FusedAdam.step() skipped a param group with no parameters before touching its step counter, so a group that is empty on one data-parallel rank and populated on another stopped counting on the empty ranks. Since step is stored in param_groups it is checkpointed, and a rank that loads its shard from a rank where the group was empty resumes with a stale step and a wrong bias correction. Move the counter update above the empty-group skip. Empty groups have no parameter to read a device from, so the capturable tensor now falls back to the device of the optimizer scratch buffer. Fixes NVIDIA#1986 Signed-off-by: Aditya Singh <adisin650@gmail.com>
github-actions
Bot
added
the
community-contribution
Greptile SummaryThe PR advances FusedAdam’s per-group step counter before skipping empty parameter groups, keeping serialized counters synchronized across distributed ranks. It also initializes capturable counters for empty groups on the optimizer scratch-buffer device and adds regression coverage for ordinary and capturable operation. Confidence Score: 5/5The PR appears safe to merge, with the empty-group counter invariant corrected and directly covered by the existing optimizer test suite. The changed ordering advances every parameter group’s counter while preserving the existing kernel skip for empty groups, and the capturable fallback uses the optimizer’s CUDA scratch-buffer device without affecting populated-group execution. Important Files Changed
Reviews (1): Last reviewed commit: "[PyTorch] Advance FusedAdam step counter..." | Re-trigger Greptile |
|
/te-ci pytorch |
|
Flagging that the three red build jobs all died before any test ran, so none of them is exercising this change.
This PR is 21 lines of Python in Not asking for anything, just did not want the three reds to read as coming from the change while |
|
Update on the run you triggered, since All four died in Build, so still nothing here has run a test. I also found the line just above the fatbinary error I quoted earlier, which points at the runner rather than at the code: The compiler was killed, so the cubin was never written. That is at 1 percent of the build, compiling |
|
Correcting my last comment. I said nothing here had run a test. That was true of the four
The red A100 is 8.0 and L40 is 8.9, so both land on 8 and fail that check. The installed I could not account for the H100 and H100 debug statuses. Their uploaded logs show no pytest failures at all, so whatever turned them red is outside the part I can read, and I did not want to guess. |
2 participants
Fixes #1986
Root cause
FusedAdam.step()opens its param-group loop with an early skip for groups that hold no parameters:The step counter is updated after that skip, so an empty group never gets one. That is harmless for a group that is empty everywhere, but a group is often empty on only some data-parallel ranks. A
no_weight_decaygroup holding just RMSNorm parameters is the usual case, and with pipeline or expert parallelism the ranks that own none of those parameters see an empty group while their peers do not.steplives inparam_groups, sostate_dict()serializes it and it goes into the checkpoint. The ranks where the group was empty writestep = nullwhile the ranks where it was populated write the true iteration count, which is exactly what the counter table in the issue shows for aPP=2, EP=4, DP=8run at iteration 2640. On resume, a rank that loads its optimizer shard from a rank where the group was empty picks up the stale value, andbias_correctioncomputes1 - beta1 ** stepfrom a step that has nothing to do with how far training actually got.Fix
Move the counter update above the empty-group skip so every group advances on every rank, then skip the kernel work for empty groups as before. This is the change suggested in the issue.
One detail the issue does not cover: with
capturable=Truethe first update createsgroup["step"]as a device tensor and takes the device fromgroup["params"][0], which an empty group does not have. The new code falls back to the device ofself._dummy_overflow_buf, the optimizer's own scratch buffer, which is allocated on CUDA in__init__. Nothing else in the loop moved, so populated groups execute exactly the same sequence as before.Verification
I do not have a GPU, so I could not run the TE test suite. Two things I did do.
The control flow itself is checked with a standalone CPU script that reproduces the loop head, before and after, on a real
torch.optim.Optimizerso thatparam_groupsandstate_dict()behave as they do in TE. The kernel launch plays no part in the defect and is omitted:Output:
The
Nonein the checkpoint onmainis thenullstep reported in the issue.The regression test in this PR is the GPU version of the same property.
test_empty_param_group_advances_stepbuilds aFusedAdamover one populated group and one empty group, steps three times, and asserts that both groups report the same step inparam_groupsand instate_dict(). It is parametrized overcapturableso the tensor-valued counter and the new device fallback are both exercised. Onmainthe test fails at the first assertion on the empty group with aKeyErrorforstep.The changed files were formatted with the repository's pinned
black24.4.2 and the pre-commit arguments, and both are unchanged by it.